#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time     : 2023/8/8 17:10 
# @Author   : 22759
# @Email    : zjh13092921608@163.com
# @Project  : threatIntelligenceProject
# @File     : sqlServerObject
# @Software : PyCharm
import json
import math
import os.path
import pymysql
import requests

from typing import Tuple, Any, Dict


class SQLServerClass(object):
    """
    数据库连接类
    :param host:     主机地址
    :param username: 用户名
    :param password: 密码
    :param database: 数据库
    :param port:     端口
    :param charset:  编码
    """
    def __init__(self, host, username, password, database, port=3306, charset='utf8'):
        self.host = host
        self.port = port
        self.cursor = None
        self.charset = charset
        self.username = username
        self.password = password
        self.database = database
        try:
            self.conn = pymysql.connect(
                host=self.host,
                user=self.username,
                password=self.password,
                database=self.database,
                port=self.port,
                charset=charset
            )
            self.cursor = self.conn.cursor()
        except pymysql.Error as error:
            print(f"无法连接服务器：{self.host}的数据库，错误信息为：{error}")

    def execute(self, sql_str, values=None) -> object:
        self.cursor.execute(sql_str, values)
        self.conn.commit()
        return self.cursor

    def fetchone(self) -> Dict[Any, None]:
        names = [desc[0] for desc in self.cursor.description]
        rows = self.cursor.fetchone()
        if rows:
            return dict(zip(names, rows))
        return {}

    def fetchall(self) -> Tuple[Tuple[Any, ...], ...]:
        return self.cursor.fetchall()

    def close(self):
        print('关闭数据库连接！！！')
        self.cursor.close()
        self.conn.close()


class RemoteRequest(object):
    """
    远程请求
    :param _type: 设备类型
    :param _ip:   请求IP地址
    :param _data: 请求数据
    """
    def __init__(self, _type, _ip, _behavior, _data=None):
        self._type, self._ip, self._behavior, self._data = _type, _ip, _behavior, _data
        self.token = self._token()

    def base_requests(self, url, method, path, filename, data=None):
        """路由中枢"""
        if self.token:
            if 'A4' in self._type:
                kwargs = {
                    'url': url,
                    'data': data,
                    'verify': False,
                    'timeout': 1200,
                    'headers': {'Authorization': self.token}
                }
                return self._request(method, path, filename, kwargs)
            elif 'T3' in self._type:
                kwargs = {
                    'url': f'{url}?auth={self.token}',
                    'data': data,
                    'verify': False,
                    'timeout': 600,
                    'headers': {"Cookie": f"sysauth={self.token}"}
                }
                return self._request(method, path, filename, kwargs)
            elif 'LAS' in self._type:
                kwargs = {
                    'url': url,
                    'data': data,
                    'verify': False,
                    'timeout': 600,
                    'headers': {"Cookie": self.token},
                }
                return self._request(method, path, filename, kwargs)
            else:
                return {'mark': False, 'message': '升级失败，操作不支持的设备'}
        else:
            return {'mark': False, 'message': '升级失败，无法获取设备权限'}

    def _token(self):
        """获取设备token令牌"""
        if 'A4' in self._type:
            response = requests.post(f'https://{self._ip}/api/api_v1/login', data=self._data, verify=False, timeout=30)
            if 200 <= response.status_code < 300:
                result = response.json()
                if result['code'] == 200:
                    return result.get('data', {}).get('Authorization', '')
        elif 'T3' in self._type:
            headers = {'Content-Type': 'application/json'}
            response = requests.get(
                url=f'https://{self._ip}/cgi-bin/webui/rpc/auth', headers=headers, data=json.dumps(self._data), verify=False, timeout=30)
            if response.status_code == 200:
                result = response.json()['result']
                return result
        elif 'LAS' in self._type:
            response = requests.get(f'https://{self._ip}/api/login_safept', params=self._data, verify=False, timeout=30)
            if 200 <= response.status_code < 300:
                resp = response.headers['set-cookie']
                return resp.split(';')[0] if resp else False
        return False

    def _request(self, method: str, path, filename, kwargs):
        """调用路由"""
        chunk_size, path = 50 * 1024 * 1024, os.path.join(path, filename)
        if os.path.exists(path):
            filesize = os.path.getsize(path)
            total = math.ceil(filesize / chunk_size)
            with open(path, 'rb') as file:
                for chunk_index in range(total):
                    """内存不够，将文件分片上传，每次上传100MB"""
                    chunk = file.read(chunk_size)
                    kwargs['data'] = {'chunk': chunk_index, 'total': total, 'behavior': self._behavior}
                    kwargs['files'] = {'file': (filename, chunk)}
                    kwargs['stream'] = True
                    if method.upper() == 'GET':
                        request = requests.get(**kwargs)
                    elif method.upper() == 'POST':
                        request = requests.post(**kwargs)
                    elif method.upper() == 'PUT':
                        request = requests.put(**kwargs)
                    else:
                        return {'mark': False, 'message': '升级失败，请求方法错误'}
                    if request.status_code not in (202, 403):
                        return {'mark': False, 'message': '升级失败，获取数据错误'}
            result = request.json()
            if request.status_code == 202:
                split_upgrade = self._split_upgrade()
                if split_upgrade['mark']:
                    return {'mark': True, 'message': '升级成功'}
                return {'mark': False, 'message': split_upgrade.get('message', '升级失败')}
            elif request.status_code == 403:
                return {'mark': False, 'message': result.get('message', '升级失败')}
            else:
                return {'mark': False, 'message': '升级失败，无法连接设备'}
        return {'mark': False, 'message': '升级失败，文件不存在'}

    def _split_upgrade(self):
        """
        升级系统
        因为系统原因导致无法直接执行升级功能，所以拆开执行，先上传文件文件上传结束在进行升级
        """
        split_system = {
            'LAS': {
                'url': f'https://{self._ip}/api/licenseInformation/shard_upgrade_new',
                'headers': {'Cookie': self.token},
                'timeout': 300,
                'verify': False,
            }
        }
        if self._type in split_system.keys():
            request = requests.post(**split_system[self._type])
            if request.status_code == 202:
                return {'mark': True, 'message': '升级成功'}
            return {'mark': False, 'message': '升级失败'}
        return {'mark': True, 'message': '不需要拆分升级的系统'}


if __name__ == '__main__':
    sql = SQLServerClass(host='10.10.0.143', username='postgres', password='isap@12345', database='isap')
    sql.execute('select * from t_device where id = 1;')
    print(sql.fetchone())
